import os.path
import time

import pandas as pd
from src.EnvironmentVariables import BASE_PATH
import numpy as np
import datetime
import matplotlib.pyplot as plt


def desc_data():
    from src.DataManager.ArticleManager import article_manager
    clause = "post_date is not null"
    column_name = ['article_id', 'title', 'post_date']
    need_pd_df = True
    _df = article_manager.get_data(clause=clause, column_name=column_name, need_pd_df=need_pd_df)
    _df['post_date'] = pd.to_datetime(_df['post_date'])
    df_order_by_date = pd.DataFrame(_df['post_date'])
    df_order_by_date.insert(1, 'times', value=1)
    df_order_by_date.set_index('post_date', inplace=True)
    df_order_by_date.resample('m').sum().to_period('m').to_csv(os.path.join(BASE_PATH, 'tmp', "month_data_num.txt"))


def dump_data():
    from src.DataManager.ArticleManager import article_manager
    clause = "post_date is not null and emotion is not null"
    column_name = ['article_id', 'post_date', 'emotion']
    need_pd_df = True
    _df = article_manager.get_data(clause=clause, column_name=column_name, need_pd_df=need_pd_df)
    _df['post_date'] = pd.to_datetime(_df['post_date'])
    _df = pd.DataFrame(_df[['post_date', 'emotion']])
    _df.insert(1, 'total', value=1)
    _df.set_index('post_date', inplace=True)
    _df.resample('d').sum().to_period('d').to_csv(os.path.join(BASE_PATH, 'data', "emotion_day_data.csv"))


def analyse_emotion():
    from src.DataManager.ArticleManager import article_manager
    from model_train import vote_ans
    batch_size = 16000
    id_offset = 3041250
    column_name = ['article_id', 'title', 'post_date']
    need_pd_df = True
    global_handle = 0
    has_handle = 1796000
    start_time = time.time()
    while True:
        clause = f" article_id>{id_offset} and post_date is not null  "
        _df = article_manager.get_data(clause=clause, column_name=column_name, need_pd_df=need_pd_df,
                                       limit=f"{batch_size}", order_by='article_id')
        # df.set_index('article_id', inplace=True)
        if len(_df) == 0:
            break
        id_offset = _df['article_id'].iloc[-1]
        ans, _ = vote_ans(_df)
        _df.insert(1, "emotion", ans)
        article_manager.update4df(_df, update_col=["emotion"])
        global_handle += len(_df)
        print(
            f"has analyse {global_handle + has_handle}, speed is {round(global_handle / (time.time() - start_time), 2)}step/s")
        # input()


def count_day_emotion_index(_df: pd.DataFrame):
    def count_emotion_index(_df: pd.DataFrame) -> float:
        if _df.total == 0:
            return np.NAN
        return np.log((1 + _df.emotion) / (1 + _df.total - _df.emotion))

    _df.insert(0, 'emotion_index', np.NAN)
    _df['emotion_index'] = _df.apply(count_emotion_index, axis=1)
    return _df


def count_ewm(_df: pd.DataFrame):
    _df['month_emotion_index'] = _df['emotion_index'].ewm(span=30, ignore_na=True).mean()
    _df['week_emotion_index'] = _df['emotion_index'].ewm(span=7, ignore_na=True).mean()
    _df['3day_emotion_index'] = _df['emotion_index'].ewm(span=3, ignore_na=True).mean()
    return _df


def plot_emotion(_df: pd.DataFrame):
    plt.figure()
    plt.title("Forum sentiment index")
    plt.xlabel('Time')  # x轴标题
    plt.ylabel('Index number')  # y轴标题
    x = _df.Date
    y0 = _df.emotion_index
    y1 = _df.month_emotion_index
    y2 = _df.week_emotion_index
    y3 = _df['3day_emotion_index']
    plt.plot(x, y0, marker='o', markersize=3)
    plt.plot(x, y1, marker='o', markersize=3)  # 绘制折线图，添加数据点，设置点的大小
    plt.plot(x, y2, marker='o', markersize=3)
    plt.plot(x, y3, marker='o', markersize=3)
    plt.legend(['Daily Emotional Index', 'Monthly Moving Average Index', 'Weekly Moving Average Index',
                '3 Day Moving Average Index'])

    plt.show()


if __name__ == '__main__':
    df = pd.read_csv(os.path.join(BASE_PATH, "data", "emotion_day_data.csv"))
    # desc_data()
    # analyse_emotion()
    # dump_data()
    df = df.rename(columns={'post_date': 'Date'})
    df['Date'] = pd.to_datetime(df['Date'])
    df.set_index('Date')
    df = df.loc[df.Date > datetime.datetime(2013, 11, 30)]
    df = count_day_emotion_index(df)
    df = count_ewm(df)

    df.to_csv(os.path.join(BASE_PATH, 'data', 'emotion_index.csv'), index=False)
    plot_emotion(df)
    # fix_empty(df)
